from .base import PoseEstimator, center_ellipse, ransac_pnp
import numpy as np
import cv2


class BA(PoseEstimator):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __name__(self):
        return "ba"

    def forward(self, C_3d, C_2d,*args, **ignore_kwargs):
        x_cnt = []
        X_cnt = []
        for c_3d, c_2d in zip(C_3d, C_2d):
            # 求解相机椭圆中心：
            x_cnt.append(center_ellipse(c_2d))
            X_cnt.append(center_ellipse(c_3d))
        # 求解相机位姿
        X_cnt = np.pad(X_cnt, ((0, 0), (0, 1)), "constant", constant_values=0)
        x_cnt = np.array(x_cnt)
        ind = ransac_pnp(self.K, X_cnt, x_cnt)
        if ind is None:
            return False, None, None
        retval, rvec, tvec = cv2.solvePnP(X_cnt[ind], x_cnt[ind], self.K, np.zeros(4))
        if not retval:
            return False, None, None
        retval, rvec, tvec = cv2.solvePnP(
            X_cnt[ind],
            x_cnt[ind],
            self.K,
            np.zeros(4),
            rvec,
            tvec,
            useExtrinsicGuess=True,
            flags=cv2.SOLVEPNP_ITERATIVE,
        )
        if not retval:
            return False, None, None
        R = cv2.Rodrigues(rvec)[0]
        return True, R, tvec
